# -*- coding: utf8 -*-

import os, sys, configparser, yara
import pymysql
import time

class YaraRuleMysql():
    def __init__(self, 
                host='127.0.0.1',
                port=3306, 
                user='root', 
                passwd='Yezishuqaz123',
                db='fhq',
                charset='utf8'):
        self._host = host
        self._port = port
        self._user = user
        self._passwd = passwd
        self._db = db
        self._charset = charset
        self.yara_rule = {}  #{(ip,port) : [ {name : [action, log, enable]} , {} ]}
        self.globe_rule = {} #[{name: level}]

    def ruleMysql(self):
        db = pymysql.connect(host=self._host,
                            port=self._port, 
                            user= self._user, 
                            passwd=self._passwd,
                            db=self._db,
                            charset=self._charset)
        cursor = db.cursor()
        #waf policy from table yara_policy
        sql = "SELECT * FROM yara_policy"
        cursor.execute(sql)
        results = cursor.fetchall()
        for row in results:
            # waf name_action_log_enable from table yara_rule_name_action
            sql_threat = "SELECT * FROM yara_rule_name_action WHERE object_index = %d"%(row[4])
            cursor.execute(sql_threat)
            results_threat = cursor.fetchall()
            value_list = [0]*cursor.rowcount
            value_index = 0
            for row_th in results_threat:
                value_list[value_index] = {str(row_th[2]):[row_th[3], row_th[4], row_th[5]]}
                value_index += 1
            self.yara_rule[(str(row[2]), str(row[3]))] = value_list
        db.close()
   
    def job(self):
        while 1:
            self.ruleMysql()
            time.sleep(60) 

    def ruleMysqlInfo(self):
        db = pymysql.connect(host=self._host,
                            port=self._port, 
                            user= self._user, 
                            passwd=self._passwd,
                            db=self._db,
                            charset=self._charset)
        cursor = db.cursor()
        #rule_info from table yara_rule
        sql = "SELECT * FROM yara_rule"
        cursor.execute(sql)
        results = cursor.fetchall()
        for row in results:
            self.globe_rule.setdefault(str(row[1]))
            self.globe_rule[str(row[1])] = row[2]
        db.close()
        return self.globe_rule

if __name__ == '__main__':
    pass
